import React, { useState } from 'react';
import { Tree } from 'antd';
import type { DataNode, TreeProps } from 'antd/es/tree';

const x = 3;
const y = 2;
const z = 1;
const defaultData: DataNode[] = [];

const generateData = (_level: number, _preKey?: React.Key, _tns?: DataNode[]) => {
  const preKey = _preKey || '0';
  const tns = _tns || defaultData;

  const children: React.Key[] = [];
  for (let i = 0; i < x; i++) {
    const key = `${preKey}-${i}`;
    tns.push({ title: key, key });
    if (i < y) {
      children.push(key);
    }
  }
  if (_level < 0) {
    return tns;
  }
  const level = _level - 1;
  children.forEach((key, index) => {
    tns[index].children = [];
    return generateData(level, key, tns[index].children);
  });
};
generateData(z);

function findNodeFromKey(data, key): [node: DataNode, i: number, nodes: DataNode[]] {
  for (let i = 0; i < data.length; i++) {
    if (data[i].key === key) {
      return [data[i], i, data];
    }
    if(data[i].children){
      const result = findNodeFromKey(data[i].children,key);
      if(result){
        return result;
      }
    }
  }
}

function removeNode(data, key): DataNode {
  const nodeObj = findNodeFromKey(data, key);
  if(nodeObj){
    const [node, index, nodes] = nodeObj;
    nodes.splice(index, 1);
    return node;
  }
}


const App: React.FC = () => {
  const [gData, setGData] = useState(defaultData);
  const [expandedKeys, setExpandedKeys] = useState(['0-0', '0-0-0', '0-0-0-0']);

  const onDragEnter: TreeProps['onDragEnter'] = (info) => {
    // console.log(info);
    // expandedKeys 需要受控时设置
    // setExpandedKeys(info.expandedKeys)
  };

  const onDrop: TreeProps['onDrop'] = (info) => {
    const dropKey = info.node.key;
    const dragKey = info.dragNode.key;
    const dropPos = info.node.pos.split('-');
    const dropPosition = info.dropPosition - Number(dropPos[dropPos.length - 1]);

    const data = [...gData];
    // Find dragObject
    const dragObj = removeNode(data, dragKey);
    const [dropObj, dropIndex, dropArr] = findNodeFromKey(data, dropKey);
    if (!info.dropToGap) {
      // Drop on the content
      dropObj.children = dropObj.children || [];
      if((info.node as any).props.expanded){
        dropObj.children.unshift(dragObj);
      }else {
        dropObj.children.push(dragObj);
      }
    }  else {
      if (dropPosition === -1) {
        dropArr.splice(dropIndex, 0, dragObj);
      } else {
        dropArr.splice(dropIndex + 1, 0, dragObj);
      }
    }
    setGData(data);
  };

  return (
    <Tree
      className="draggable-tree"
      defaultExpandedKeys={expandedKeys}
      draggable
      blockNode
      onDragEnter={onDragEnter}
      onDrop={onDrop}
      treeData={gData}
    />
  );
};

export default App;